from peewee import *
import json
from datetime import datetime
from typing import Dict, List, Optional, Any
import logging
from config.config import Config
from exception.base_exceptions import DataProcessingException
from exception.exception_handler import global_exception_handler

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库连接
database = MySQLDatabase(
    Config.database, 
    host=Config.host, 
    user=Config.user, 
    password=Config.password,
    charset='utf8mb4')

class BrainWave(Model):
    """
    Brain Wave数据模型类
    基于Peewee ORM的脑电波数据模型
    """
    
    id = CharField(primary_key=True, max_length=50)
    session_id = CharField(unique=True, max_length=100, null=False)
    room_id = CharField(max_length=100, default='default_room')
    created_at = TimestampField(default=datetime.now)
    total_seconds = IntegerField(null=False)
    
    # JSON字段
    wave = TextField(null=True)  # 原始脑电数据
    delta_wave = TextField(null=True)  # Delta波数据
    theta_wave = TextField(null=True)  # Theta波数据
    alpha_wave = TextField(null=True)  # Alpha波数据
    beta_wave = TextField(null=True)  # Beta波数据
    gamma_wave = TextField(null=True)  # Gamma波数据
    metadata = TextField(null=True)  # 元数据
    
    class Meta:
        database = database
        table_name = 'brain_wave'
    
    def get_wave_data(self) -> Dict[str, Any]:
        """获取原始脑电数据"""
        if self.wave:
            return json.loads(self.wave)
        return {}
    
    def set_wave_data(self, data: Dict[str, Any]):
        """设置原始脑电数据"""
        self.wave = json.dumps(data, ensure_ascii=False)
    
    def get_delta_wave_data(self) -> Dict[str, Any]:
        """获取Delta波数据"""
        if self.delta_wave:
            return json.loads(self.delta_wave)
        return {}
    
    def set_delta_wave_data(self, data: Dict[str, Any]):
        """设置Delta波数据"""
        self.delta_wave = json.dumps(data, ensure_ascii=False)
    
    def get_theta_wave_data(self) -> Dict[str, Any]:
        """获取Theta波数据"""
        if self.theta_wave:
            return json.loads(self.theta_wave)
        return {}
    
    def set_theta_wave_data(self, data: Dict[str, Any]):
        """设置Theta波数据"""
        self.theta_wave = json.dumps(data, ensure_ascii=False)
    
    def get_alpha_wave_data(self) -> Dict[str, Any]:
        """获取Alpha波数据"""
        if self.alpha_wave:
            return json.loads(self.alpha_wave)
        return {}
    
    def set_alpha_wave_data(self, data: Dict[str, Any]):
        """设置Alpha波数据"""
        self.alpha_wave = json.dumps(data, ensure_ascii=False)
    
    def get_beta_wave_data(self) -> Dict[str, Any]:
        """获取Beta波数据"""
        if self.beta_wave:
            return json.loads(self.beta_wave)
        return {}
    
    def set_beta_wave_data(self, data: Dict[str, Any]):
        """设置Beta波数据"""
        self.beta_wave = json.dumps(data, ensure_ascii=False)
    
    def get_gamma_wave_data(self) -> Dict[str, Any]:
        """获取Gamma波数据"""
        if self.gamma_wave:
            return json.loads(self.gamma_wave)
        return {}
    
    def set_gamma_wave_data(self, data: Dict[str, Any]):
        """设置Gamma波数据"""
        self.gamma_wave = json.dumps(data, ensure_ascii=False)
    
    def get_metadata(self) -> Dict[str, Any]:
        """获取元数据"""
        if self.metadata:
            return json.loads(self.metadata)
        return {}
    
    def set_metadata(self, data: Dict[str, Any]):
        """设置元数据"""
        self.metadata = json.dumps(data, ensure_ascii=False)
    
    @classmethod
    def create_session(cls, session_data: Dict[str, Any]) -> 'BrainWave':
        """
        创建睡眠会话
        
        Args:
            session_data: 会话数据字典
            
        Returns:
            BrainWave: 创建的会话实例
        """
        try:
            session = cls.create(
                id=session_data.get('id'),
                session_id=session_data.get('session_id'),
                room_id=session_data.get('room_id', 'default_room'),
                total_seconds=session_data.get('total_seconds')
            )
            
            # 设置JSON数据
            if 'wave' in session_data:
                session.set_wave_data(session_data['wave'])
            if 'delta_wave' in session_data:
                session.set_delta_wave_data(session_data['delta_wave'])
            if 'theta_wave' in session_data:
                session.set_theta_wave_data(session_data['theta_wave'])
            if 'alpha_wave' in session_data:
                session.set_alpha_wave_data(session_data['alpha_wave'])
            if 'beta_wave' in session_data:
                session.set_beta_wave_data(session_data['beta_wave'])
            if 'gamma_wave' in session_data:
                session.set_gamma_wave_data(session_data['gamma_wave'])
            if 'metadata' in session_data:
                session.set_metadata(session_data['metadata'])
            
            session.save()
            logger.info(f"睡眠会话创建成功: {session.session_id}")
            return session
            
        except Exception as e:
            exception = DataProcessingException(f"创建睡眠会话失败，会话ID: {session_data.get('session_id', 'unknown')}，错误详情: {str(e)}")
            global_exception_handler.handle_exception(exception)
            raise
    
    @classmethod
    def get_by_session_id(cls, session_id: str) -> Optional['BrainWave']:
        """
        根据会话ID获取会话
        
        Args:
            session_id: 会话ID
            
        Returns:
            BrainWave: 会话实例，如果不存在返回None
        """
        try:
            return cls.get(cls.session_id == session_id)
        except cls.DoesNotExist:
            logger.warning(f"未找到会话: {session_id}")
            return None
        except Exception as e:
            exception = DataProcessingException(f"获取会话失败，会话ID: {session_id}，错误详情: {str(e)}")
            global_exception_handler.handle_exception(exception)
            return None
    
    @classmethod
    def get_all_sessions(cls, limit: int = 10) -> List['BrainWave']:
        """
        获取所有会话列表
        
        Args:
            limit: 返回记录数量限制
            
        Returns:
            List[BrainWave]: 会话列表
        """
        try:
            return list(cls.select().order_by(cls.created_at.desc()).limit(limit))
        except Exception as e:
            exception = DataProcessingException(f"获取会话列表失败，限制数量: {limit}，错误详情: {str(e)}")
            global_exception_handler.handle_exception(exception)
            return []
    
    @classmethod
    def get_sessions_by_date_range(cls, start_date: datetime, end_date: datetime) -> List['BrainWave']:
        """
        根据日期范围获取会话
        
        Args:
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            List[BrainWave]: 会话列表
        """
        try:
            return list(cls.select().where(
                (cls.created_at >= start_date) & (cls.created_at <= end_date)
            ).order_by(cls.created_at.desc()))
        except Exception as e:
            exception = DataProcessingException(f"根据日期范围获取会话失败，开始时间: {start_date}，结束时间: {end_date}，错误详情: {str(e)}")
            global_exception_handler.handle_exception(exception)
            return []
    
    @classmethod
    def get_statistics(cls) -> Dict[str, Any]:
        """
        获取统计信息
        
        Returns:
            Dict: 统计信息
        """
        try:
            total_sessions = cls.select().count()
            total_seconds = cls.select(fn.SUM(cls.total_seconds)).scalar() or 0
            avg_seconds = cls.select(fn.AVG(cls.total_seconds)).scalar() or 0
            
            first_session = cls.select().order_by(cls.created_at.asc()).first()
            last_session = cls.select().order_by(cls.created_at.desc()).first()
            
            return {
                'total_sessions': total_sessions,
                'total_sleep_seconds': total_seconds,
                'avg_sleep_seconds': round(avg_seconds, 2),
                'total_sleep_hours': round(total_seconds / 3600, 2),
                'avg_sleep_hours': round(avg_seconds / 3600, 2),
                'first_session': first_session.created_at if first_session else None,
                'last_session': last_session.created_at if last_session else None
            }
        except Exception as e:
            exception = DataProcessingException(f"获取统计信息失败，错误详情: {str(e)}")
            global_exception_handler.handle_exception(exception)
            return {}


# 数据库连接管理
def connect_db():
    """连接数据库"""
    try:
        database.connect()
        logger.info("数据库连接成功")
    except Exception as e:
        exception = DataProcessingException(f"数据库连接失败，主机: {Config.host}，数据库: {Config.database}，错误详情: {str(e)}")
        global_exception_handler.handle_exception(exception)
        raise

def close_db():
    """关闭数据库连接"""
    try:
        if not database.is_closed():
            database.close()
        logger.info("数据库连接已关闭")
    except Exception as e:
        exception = DataProcessingException(f"关闭数据库连接失败，错误详情: {str(e)}")
        global_exception_handler.handle_exception(exception)

def create_tables():
    """创建数据表"""
    try:
        database.create_tables([BrainWave])
        logger.info("数据表创建成功")
    except Exception as e:
        exception = DataProcessingException(f"创建数据表失败，表名: brain_wave，错误详情: {str(e)}")
        global_exception_handler.handle_exception(exception)
        raise


# 使用示例
if __name__ == "__main__":
    # 连接数据库
    connect_db()
    
    try:
        # 创建表
        create_tables()
        
        # 测试数据
        test_session_data = {
            'id': 'test_001',
            'session_id': 'sleep_session_001',
            'total_seconds': 7200,
            'wave': {
                'total_seconds': 7200,
                'frame_rate': 128,
                'frames': [[1, 12324], [2, 32223], [3, 15678]]
            },
            'delta_wave': {
                'total_seconds': 7200,
                'frame_rate': 1,
                'frames': [[1, 23], [2, 25], [3, 28]]
            },
            'metadata': {
                'device_info': {
                    'device_name': 'NeuroSky MindWave',
                    'device_type': 'EEG'
                },
                'signal_quality': 'good'
            }
        }
        
        # 创建会话
        session = BrainWave.create_session(test_session_data)
        print(f"创建会话成功: {session.session_id}")
        
        # 查询会话
        found_session = BrainWave.get_by_session_id('sleep_session_001')
        if found_session:
            print(f"查询到会话: {found_session.session_id}")
            print(f"Delta波数据: {found_session.get_delta_wave_data()}")
        
        # 获取统计信息
        stats = BrainWave.get_statistics()
        print(f"统计信息: {stats}")
        
    finally:
        # 关闭数据库连接
        close_db()